
import os

import dotenv
from langchain_core.output_parsers import StrOutputParser, JsonOutputParser
from langchain_core.prompts import PromptTemplate
from langchain_openai import ChatOpenAI

#字符串输出解析器 StrOutputParser
def fun_1():
    dotenv.load_dotenv()

    chat_model = ChatOpenAI(
        model=os.getenv("MODEL_NAME")
        , base_url=os.getenv("BASE_URL")
        , api_key=os.getenv("API_KEY")
        , temperature=0.8  # 精确度
        , max_tokens=1024  # 回答长度
        , streaming=True  # 是否流式返回
    )

    response = chat_model.invoke("介绍下LLM")

    #获取字符串输出结果
    #print(response.content) #方式1，使用content属性
    #方式2，使用StrOutputParser
    parser = StrOutputParser()
    response_str = parser.invoke(response)
    print(response_str)

#fun_1()

#Json输出解析器 方式1
def fun_2():
    dotenv.load_dotenv()
    chat_model = ChatOpenAI(
        model=os.getenv("MODEL_NAME")
        , base_url=os.getenv("BASE_URL")
        , api_key=os.getenv("API_KEY")
        , temperature=0.8  # 精确度
        , max_tokens=1024  # 回答长度
        , streaming=True  # 是否流式返回
    )
    parser = JsonOutputParser()
    response = chat_model.invoke("介绍下LLM,问题用Q表示，答案用A表示，请用json格式返回结果")
    response_json = parser.invoke(response)
    print(response_json)

#fun_2()

#Json输出解析器 方式2
def fun_3():
    dotenv.load_dotenv()
    chat_model = ChatOpenAI(
        model=os.getenv("MODEL_NAME")
        , base_url=os.getenv("BASE_URL")
        , api_key=os.getenv("API_KEY")
        , temperature=0.8  # 精确度
        , max_tokens=1024  # 回答长度
    )

    question = "介绍下PYTHON"
    parser = JsonOutputParser()

    prompt_template = PromptTemplate.from_template(
        template="回答用户的问题\n 满足的格式为：{format_instructions}\n 问题是：{question}\n"
        , partial_variables={"format_instructions": parser.get_format_instructions()}
    )
    prompt = prompt_template.invoke({"question": question})
    response = chat_model.invoke(prompt)
    response_json = parser.invoke(response)
    print(response_json)

#fun_3()

#使用管道
def fun_4():
    dotenv.load_dotenv()
    chat_model = ChatOpenAI(
        model=os.getenv("MODEL_NAME")
        , base_url=os.getenv("BASE_URL")
        , api_key=os.getenv("API_KEY")
        , temperature=0.8  # 精确度
        , max_tokens=1024  # 回答长度
        , streaming=True  # 是否流式返回
    )

    question = "介绍下PYTHON"
    parser = JsonOutputParser()

    prompt_template = PromptTemplate.from_template(
        template="回答用户的问题\n 满足的格式为：{format_instructions}\n 问题是：{question}\n"
        , partial_variables={"format_instructions": parser.get_format_instructions()}
    )
    char = prompt_template | chat_model | parser
    response_json = char.invoke({"question": question})
    print(response_json)

#fun_4()